# 1. how to use G?
import matplotlib.pyplot as plt
import torch
import torchvision.utils as vutil

from config import HP
from dataset_face import invTrans
from generator import Generator

# new a generator model instance
G = Generator()
checkpoint = torch.load('./model_save/model_g_34_49000.pth', map_location='cpu')
G.load_state_dict(checkpoint['model_state_dict'])
G.to(HP.device)
G.eval()  # set evaluation mode

while 1:
    # 1. Disentangled representation: manual set Z： [0.3, 0, ]
    # 2. any input: z: fuzzy image -> high resolution image / mel -> audio/speech(decoder)

    latent_z = torch.randn(size=(HP.batch_size, HP.z_dim), device=HP.device)
    fake_faces = G(latent_z)
    grid = vutil.make_grid(fake_faces, nrow=12)  # format into a "big" image
    plt.title('generated 12 x 12 faces')
    plt.imshow(invTrans(grid).permute(1, 2, 0))  # HWC
    plt.show()
    print("12 x 12 faces generated! ")
    input()
